package com.jy.oms.framework.thread;

import org.apache.commons.collections4.MapUtils;
import org.apache.commons.lang3.tuple.MutablePair;
import org.apache.commons.lang3.tuple.Pair;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationListener;
import org.springframework.context.event.ContextRefreshedEvent;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.stereotype.Component;

import java.util.LinkedHashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.RejectedExecutionHandler;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 线程池管理服务
 *
 * @author liuquan
 */
@Component
public class ThreadPoolManager implements ApplicationListener<ContextRefreshedEvent> {
    private static final String springThreadPoolTaskKey = "spring-threadPoolTask-";
    private static final String springThreadPoolKey = "spring-threadPool-";
    private AtomicInteger springThreadPoolId = new AtomicInteger(0);
    private Map<String, Pair<String, ThreadPoolTaskExecutor>> springThreadPoolTaskExecutorMap = new ConcurrentHashMap<>();
    private Map<String, Pair<String, ThreadPoolExecutor>> springThreadPoolExecutorMap = new ConcurrentHashMap<>();

    @Override
    public void onApplicationEvent(ContextRefreshedEvent event) {
        ApplicationContext applicationContext = event.getApplicationContext();
        if (applicationContext.getParent() == null) {
            Map<String, ThreadPoolTaskExecutor> threadPoolTaskExecutorMap = applicationContext.getBeansOfType(ThreadPoolTaskExecutor.class);
            if (MapUtils.isNotEmpty(threadPoolTaskExecutorMap)) {
                for (Map.Entry<String, ThreadPoolTaskExecutor> entry : threadPoolTaskExecutorMap.entrySet()) {
                    springThreadPoolTaskExecutorMap.put(springThreadPoolTaskKey + springThreadPoolId.incrementAndGet(),
                            new MutablePair<>(entry.getKey(), entry.getValue()));
                }
            }

            Map<String, ThreadPoolExecutor> threadPoolExecutorMap = applicationContext.getBeansOfType(ThreadPoolExecutor.class);
            if (MapUtils.isNotEmpty(threadPoolExecutorMap)) {
                for (Map.Entry<String, ThreadPoolExecutor> entry : threadPoolExecutorMap.entrySet()) {
                    springThreadPoolExecutorMap.put(springThreadPoolKey + springThreadPoolId.incrementAndGet(),
                            new MutablePair<>(entry.getKey(), entry.getValue()));
                }
            }
        }
    }

    /**
     * 动态修改线程池参数
     *
     * @param threadPoolKey   线程池key
     * @param corePoolSize    核心线程池数
     * @param maximumPoolSize 最大线程数
     */
    public void configureThreadPool(String threadPoolKey,
                                    Integer corePoolSize,
                                    Integer maximumPoolSize) {
        CustomThreadPoolExecutor customThreadPoolExecutor = CustomThreadPoolManager.getInstance().getCustomThreadPoolExecutorMap().get(threadPoolKey);
        if (customThreadPoolExecutor != null) {
            customThreadPoolExecutor.setCorePoolSize(corePoolSize);
            customThreadPoolExecutor.setMaximumPoolSize(maximumPoolSize);
            return;
        }

        Pair<String, ThreadPoolTaskExecutor> springThreadPoolTaskExecutorPair = springThreadPoolTaskExecutorMap.get(threadPoolKey);
        if (springThreadPoolTaskExecutorPair != null) {
            ThreadPoolTaskExecutor threadPoolTaskExecutor = springThreadPoolTaskExecutorPair.getRight();
            threadPoolTaskExecutor.setCorePoolSize(corePoolSize);
            threadPoolTaskExecutor.setMaxPoolSize(maximumPoolSize);
            return;
        }

        Pair<String, ThreadPoolExecutor> springThreadPoolExecutorPair = springThreadPoolExecutorMap.get(threadPoolKey);
        if (springThreadPoolExecutorPair != null) {
            ThreadPoolExecutor threadPoolExecutor = springThreadPoolExecutorPair.getRight();
            threadPoolExecutor.setCorePoolSize(corePoolSize);
            threadPoolExecutor.setMaximumPoolSize(maximumPoolSize);
            return;
        }
    }

    /**
     * 获取线程池状态
     */
    public Map<String, ThreadPoolStatus> getThreadPoolStatus() {
        Map<String, ThreadPoolStatus> threadPoolStatusMap = new LinkedHashMap<>();
        if (MapUtils.isNotEmpty(CustomThreadPoolManager.getInstance().getCustomThreadPoolExecutorMap())) {
            for (Map.Entry<String, CustomThreadPoolExecutor> entry : CustomThreadPoolManager.getInstance().getCustomThreadPoolExecutorMap().entrySet()) {
                CustomThreadPoolExecutor threadPoolExecutor = entry.getValue();
                ThreadPoolStatus threadPoolStatus = new ThreadPoolStatus();
                threadPoolStatus.setThreadPoolId(entry.getKey());
                threadPoolStatus.setThreadPoolName(threadPoolExecutor.getThreadPoolName());
                threadPoolStatus.setCorePoolSize(threadPoolExecutor.getCorePoolSize());
                threadPoolStatus.setMaximumPoolSize(threadPoolExecutor.getMaximumPoolSize());
                threadPoolStatus.setPoolSize(threadPoolExecutor.getPoolSize());
                threadPoolStatus.setLargestPoolSize(threadPoolExecutor.getLargestPoolSize());
                threadPoolStatus.setTaskCount(threadPoolExecutor.getTaskCount());
                threadPoolStatus.setActiveCount(threadPoolExecutor.getActiveCount());
                threadPoolStatus.setCompletedTaskCount(threadPoolExecutor.getCompletedTaskCount());
                threadPoolStatus.setQueueSize(threadPoolExecutor.getQueue().size());
                threadPoolStatus.setReject(this.getReject(threadPoolExecutor.getRejectedExecutionHandler()));

                threadPoolStatusMap.put(entry.getKey(), threadPoolStatus);
            }
        }

        if (MapUtils.isNotEmpty(springThreadPoolTaskExecutorMap)) {
            for (Map.Entry<String, Pair<String, ThreadPoolTaskExecutor>> entry : springThreadPoolTaskExecutorMap.entrySet()) {
                String beanName = entry.getValue().getLeft();
                ThreadPoolTaskExecutor threadPoolTaskExecutor = entry.getValue().getRight();
                ThreadPoolStatus threadPoolStatus = new ThreadPoolStatus();
                threadPoolStatus.setThreadPoolId(entry.getKey());
                threadPoolStatus.setThreadPoolName(beanName);
                threadPoolStatus.setCorePoolSize(threadPoolTaskExecutor.getCorePoolSize());
                threadPoolStatus.setMaximumPoolSize(threadPoolTaskExecutor.getMaxPoolSize());
                threadPoolStatus.setPoolSize(threadPoolTaskExecutor.getPoolSize());
                threadPoolStatus.setLargestPoolSize(threadPoolTaskExecutor.getThreadPoolExecutor().getLargestPoolSize());
                threadPoolStatus.setTaskCount(threadPoolTaskExecutor.getThreadPoolExecutor().getTaskCount());
                threadPoolStatus.setActiveCount(threadPoolTaskExecutor.getActiveCount());
                threadPoolStatus.setCompletedTaskCount(threadPoolTaskExecutor.getThreadPoolExecutor().getCompletedTaskCount());
                threadPoolStatus.setQueueSize(threadPoolTaskExecutor.getThreadPoolExecutor().getQueue().size());
                threadPoolStatus.setReject(this.getReject(threadPoolTaskExecutor.getThreadPoolExecutor().getRejectedExecutionHandler()));

                threadPoolStatusMap.put(entry.getKey(), threadPoolStatus);
            }
        }

        if (MapUtils.isNotEmpty(springThreadPoolExecutorMap)) {
            for (Map.Entry<String, Pair<String, ThreadPoolExecutor>> entry : springThreadPoolExecutorMap.entrySet()) {
                String beanName = entry.getValue().getLeft();
                ThreadPoolExecutor threadPoolExecutor = entry.getValue().getRight();
                ThreadPoolStatus threadPoolStatus = new ThreadPoolStatus();
                threadPoolStatus.setThreadPoolId(entry.getKey());
                threadPoolStatus.setThreadPoolName(beanName);
                threadPoolStatus.setCorePoolSize(threadPoolExecutor.getCorePoolSize());
                threadPoolStatus.setMaximumPoolSize(threadPoolExecutor.getMaximumPoolSize());
                threadPoolStatus.setPoolSize(threadPoolExecutor.getPoolSize());
                threadPoolStatus.setLargestPoolSize(threadPoolExecutor.getLargestPoolSize());
                threadPoolStatus.setTaskCount(threadPoolExecutor.getTaskCount());
                threadPoolStatus.setActiveCount(threadPoolExecutor.getActiveCount());
                threadPoolStatus.setCompletedTaskCount(threadPoolExecutor.getCompletedTaskCount());
                threadPoolStatus.setQueueSize(threadPoolExecutor.getQueue().size());
                threadPoolStatus.setReject(this.getReject(threadPoolExecutor.getRejectedExecutionHandler()));

                threadPoolStatusMap.put(entry.getKey(), threadPoolStatus);
            }
        }

        return threadPoolStatusMap;
    }

    private String getReject(RejectedExecutionHandler handler) {
        if (handler == null) {
            return "";
        }

        if (handler instanceof ThreadPoolExecutor.AbortPolicy) {
            return "Abort";
        } else if (handler instanceof ThreadPoolExecutor.CallerRunsPolicy) {
            return "CallerRun";
        } else if (handler instanceof ThreadPoolExecutor.DiscardOldestPolicy) {
            return "DiscardOldest";
        } else if (handler instanceof ThreadPoolExecutor.DiscardPolicy) {
            return "Discard";
        } else {
            return "自定义";
        }
    }
}
